#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/3/12 17:32
# @Author  : LiShan
# @Email   : lishan_1997@126.com
# @File    : draw_image.py
# @Note    : this is note
import matplotlib.pyplot as plt
import numpy as np

# 支持中文
plt.rcParams['font.sans-serif'] = ['SimSun']  # 宋体中文
plt.rcParams['axes.unicode_minus'] = False    # 正常显示负号


# 读取excel数据
def excel_read(path, row):
    import xlrd
    import datetime
    work_book = xlrd.open_workbook(path)
    sheets_name = work_book.sheet_names()
    for index in range(work_book.nsheets):
        sheet = work_book.sheet_by_name(sheets_name[index])
        sheets_name[index] = []
        # 按行读取
        if row:
            for row in range(sheet.nrows):
                line = []
                for col in range(0, sheet.ncols):
                    """ 0 empty, 1 string, 2 number, 3 date, 4 boolean, 5 error """
                    if sheet.cell(row, col).ctype == 3:
                        # 保存为时间格式
                        date = xlrd.xldate_as_tuple(sheet.cell(row, col).value, 0)
                        line.append(str(datetime.time(*date[3::])))
                    else:
                        line.append(sheet.cell_value(row, col))
                sheets_name[index].append(line)
        # 按列读取
        else:
            for col in range(sheet.ncols):
                line = []
                for row in range(0, sheet.nrows):
                    """ 0 empty, 1 string, 2 number, 3 date, 4 boolean, 5 error """
                    if sheet.cell(row, col).ctype == 3:
                        # 保存为时间格式
                        date = xlrd.xldate_as_tuple(sheet.cell(row, col).value, 0)
                        line.append(str(datetime.time(*date[3::])))
                    else:
                        line.append(sheet.cell_value(row, col))
                sheets_name[index].append(line)
    return sheets_name


# 随机生成n种RGB颜色
def colourlist_generator(n):
    import random
    rangelist = ['1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F']
    colours = []
    for _ in range(int(n)):
        colours.append("#" + "".join([random.choice(rangelist) for _ in range(6)]))
    return colours


"""5种模型与预测值与真实值差距堆叠柱状图"""
fig, ax = plt.subplots()

# 准备数据
data = []
# 按列读取数据
excel_data = excel_read("./5models.xlsx", row=False)
print("读取文件完成,按列读取5models.xlsx文件......")
for i, value in enumerate(excel_data[0]):
    # 弹出真实参考值
    ref_value = value.pop(0)
    data.append(list(abs(np.array(ref_value) - np.array(value))))
# 弹出最后的求和值
data.pop()

# 颜色表
color_table = colourlist_generator(len(data))

# 绘图
updata = np.array(0)
for i in range(len(data)):
    plt.bar(range(len(data[i])), data[i], color=color_table[i], width=0.6, bottom=updata, label="id_" + str(i + 1))
    updata = list(np.array(data[i]) + np.array(updata))

# 设置坐标轴刻度
x_ticks = ["TGG", "GCN+LSTM", "GCN(优化后)", "GRU", "LSTM"]
plt.xticks(np.arange(len(x_ticks)), x_ticks)
plt.yticks([])

# 设置数据标注
for a, b in zip(range(len(updata)), updata):
    plt.text(a, b + 0.5, '%.0f' % b, ha='center')

# 去除边框
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['bottom'].set_visible(False)
ax.spines['left'].set_visible(False)

# 显示图例
box = ax.get_position()
ax.set_position([box.x0, box.y0, box.width * 0.9, box.height])
font = {'family': 'Times New Roman', 'weight': 'normal', 'size': 6.0}
ax.legend(bbox_to_anchor=(1, 0), loc='lower left', ncol=2, columnspacing=0.1,
          labelspacing=0.1, prop=font, frameon=False)

# 保存图像
plt.savefig('./model.svg')
print("保存完成，5种模型与预测值与真实值差距堆叠柱状图......")

# 显示图像
plt.show()

"""5种模型预测结果与真实值对比折线图"""

# 准备数据
data = excel_read("./5models.xlsx", row=True)[0]
print("读取文件完成,按行读取5models.xlsx文件......")
model = ["TGG", "GCN+LSTM", "GCN(优化后)", "GRU", "LSTM"]

for i in range(len(model)):
    # 绘图
    plt.plot(data[0][:-1], color="black", marker='1', linewidth=1, markersize=5)
    plt.plot(data[i + 1][:-1], color="gray", marker='2', linewidth=1, markersize=5)

    # 设置坐标轴刻度
    plt.xticks(range(0, 70, 10))
    plt.yticks(range(0, 300, 50))

    # 设置坐标轴名称
    plt.xlabel('NODE_ID', fontproperties='Times New Roman', size=10.5)
    plt.ylabel('VALUE', fontproperties='Times New Roman', size=10.5)

    # 设置图例
    plt.legend(("REAL_VALUE", "PREDICT_VALUE"), loc="upper left", frameon=False)

    # 设置标题
    plt.title(model[i] + '模型预测结果与真实值对比', fontproperties='SimSun', size=10.5)
    # fontproperties='Times New Roman', size=10.5
    # fontproperties='宋体', size=10.5

    # 保存图像
    plt.savefig('./model_' + model[i] + '.svg')
    print("保存完成," + model[i] + "模型与预测值与真实值对比折线图......")

    # 显示图像
    plt.show()


"""损失数据对比图"""

# 准备数据
loss_data = excel_read("./loss_yes.xlsx", row=True)[0][:-1]
x = np.linspace(0, 200, 20)

# 绘图
plt.plot(x, loss_data[0], color="black", linestyle='--', marker='*', linewidth=1, markersize=5)
plt.plot(x, loss_data[1], color="black", linestyle='-', marker='^', linewidth=1, markersize=5)
plt.plot(x, loss_data[2], color="black", linestyle=':', marker='p', linewidth=1, markersize=5)
plt.plot(x, loss_data[3], color="black", linestyle='-', marker='d', linewidth=1, markersize=5)
plt.plot(x, loss_data[4], color="black", linestyle='--', marker='o', linewidth=1, markersize=5)

# 设置坐标轴范围
plt.xlim([0, 200])
plt.ylim([0.000, 0.040])

# 设置坐标轴刻度
plt.xticks(range(0, 225, 25))
plt.yticks(np.arange(0.000, 0.045, 0.005))

# 设置坐标轴名称
plt.xlabel('EPOCHS', fontproperties='Times New Roman', size=10.5)
plt.ylabel('LOSS', fontproperties='Times New Roman', size=10.5)

# 设置网格
plt.grid()

# 设置图例
plt.legend(("TGG", "GCN+LSTM", "GCN(优化后)", "GRU", "LSTM"), loc="upper right", frameon=False)

# 设置标题
plt.title('5种模型损失变化情况', fontproperties='SimSun', size=10.5)

# 保存图像
plt.savefig('./compar_loss.svg')
print("保存完成,模型损失率对比折线图......")

# 显示图像
plt.show()


"""测试、验证数据对比"""
# 准备数据
train_data = excel_read("./train_acc.xlsx", row=True)[0][:-1]
print("读取文件完成,按行读取train_acc.xlsx文件......")
val_data = excel_read("./val_acc.xlsx", row=True)[0][:-1]
print("读取文件完成,按行读取val_acc.xlsx文件......")

acc = [train_data, val_data]
y_label = ['ACC_TRAIN', 'ACC_VAL']
legend_lable = ["TGG", "GCN+LSTM", "GCN(优化后)", "GRU", "LSTM"]
title_label = ['测试集上5种模型准确率变化情况', '验证集上5种模型准确率变化情况']
save_name = ['./compar_train.svg', './compar_val.svg']

for i in range(len(acc)):
    # 绘图
    plt.plot(x, acc[i][0], color="black", linestyle='--', marker='*', linewidth=1, markersize=5)
    plt.plot(x, acc[i][1], color="black", linestyle='-', marker='^', linewidth=1, markersize=5)
    plt.plot(x, acc[i][2], color="black", linestyle=':', marker='p', linewidth=1, markersize=5)
    plt.plot(x, acc[i][3], color="black", linestyle='-', marker='d', linewidth=1, markersize=5)
    plt.plot(x, acc[i][4], color="black", linestyle='--', marker='o', linewidth=1, markersize=5)

    # 设置坐标轴范围
    plt.xlim([0, 200])
    plt.ylim([0.3, 1.0])

    # 设置坐标轴刻度
    plt.xticks(range(0, 225, 25))
    plt.yticks(np.arange(0.3, 1.1, 0.1))

    # 设置坐标轴名称
    plt.xlabel('EPOCHS', fontproperties='Times New Roman', size=10.5)
    plt.ylabel(y_label[i], fontproperties='Times New Roman', size=10.5)

    # 设置网格
    plt.grid()

    # 设置图例
    plt.legend(legend_lable, loc="lower right", frameon=False)

    # 设置标题
    plt.title(title_label[i], fontproperties='SimSun', size=10.5)

    # 保存图像
    plt.savefig(save_name[i])
    print("保存完成,模型准确率对比折线图......")

    # 显示图像
    plt.show()
